#![allow(dead_code)]

use anyhow::{Ok, Result};
use clap::Parser;
use hyper_proxy::Intercept;
use log::*;
use mitm_core::{CertificateAuthority, Proxy};
use rule::RuleHttpHandler;
use rustls_pemfile as pemfile;
use std::{fs, sync::Arc};

// 条件导入被动扫描相关模块
#[cfg(feature = "passive-scan")]
use good_mitm::scan_handler::ScanHttpHandler;
#[cfg(feature = "passive-scan")]
use passive_scan::PassiveScanEngine;

// 条件导入主动扫描相关模块
#[cfg(feature = "active-scan")]
use active_scan::{ActiveScanConfig, ActiveScanEngine, MatchTarget, RegexMatcher, RegexRule, RuleMatcher, Severity};

use good_mitm::*;

#[derive(Parser)]
#[clap(name = "Good Man in the Middle", version, about, author)]
struct AppOpts {
    #[clap(subcommand)]
    subcmd: SubCommand,
}

#[derive(Parser)]
enum SubCommand {
    /// run proxy serve
    Run(Run),
    /// run proxy with passive scanning
    RunScan(RunScan),
    /// gen your own ca cert and private key
    Genca(Genca),
    /// start passive scan server
    Scan(Scan),
    /// show scan results
    ShowResults(ShowResults),
    /// start active scan (crawler-based)
    ActiveScan(ActiveScan),
}

#[derive(Parser)]
struct Run {
    #[clap(
        short,
        long,
        default_value = "ca/private.key",
        help = "private key file path"
    )]
    key: String,
    #[clap(short, long, default_value = "ca/cert.crt", help = "cert file path")]
    cert: String,
    #[clap(short, long, help = "load rules from file or dir")]
    rule: String,
    #[clap(short, long, default_value = "127.0.0.1:34567", help = "bind address")]
    bind: String,
    #[clap(short, long, help = "upstream proxy")]
    proxy: Option<String>,
}

#[derive(Parser)]
struct RunScan {
    #[clap(
        short,
        long,
        default_value = "ca/private.key",
        help = "private key file path"
    )]
    key: String,
    #[clap(short, long, default_value = "ca/cert.crt", help = "cert file path")]
    cert: String,
    #[clap(short, long, help = "load rules from file or dir")]
    rule: String,
    #[clap(short, long, default_value = "127.0.0.1:34567", help = "bind address")]
    bind: String,
    #[clap(short, long, help = "upstream proxy")]
    proxy: Option<String>,
    #[clap(long, help = "scan config file")]
    scan_config: Option<String>,
    #[clap(long, help = "enable debug mode")]
    debug: bool,
}

#[derive(Parser)]
struct Genca {
    #[clap(short, long, help = "install cert on your trust zone")]
    trust: bool,
}

#[derive(Parser)]
struct Scan {
    #[clap(short, long, help = "scan config file")]
    config: Option<String>,
    #[clap(long, help = "scan rules directory")]
    rules_dir: Option<String>,
    #[clap(long, help = "output directory for results")]
    output: Option<String>,
    #[clap(long, help = "enable debug mode")]
    debug: bool,
}

#[derive(Parser)]
struct ShowResults {
    #[clap(short, long, help = "results database path")]
    database: Option<String>,
    #[clap(long, help = "filter by severity")]
    severity: Option<String>,
    #[clap(long, help = "limit number of results")]
    limit: Option<usize>,
    #[clap(short, long, help = "output format (json, table)")]
    format: Option<String>,
}

#[derive(Parser)]
struct ActiveScan {
    #[clap(help = "target URL to scan")]
    target: String,
    #[clap(short, long, help = "active scan config file")]
    config: Option<String>,
    #[clap(long, help = "scan rules directory")]
    rules_dir: Option<String>,
    #[clap(short, long, help = "output directory for results")]
    output: Option<String>,
    #[clap(long, help = "max crawl depth", default_value = "3")]
    max_depth: usize,
    #[clap(long, help = "max pages to crawl", default_value = "1000")]
    max_pages: usize,
    #[clap(long, help = "concurrent requests", default_value = "10")]
    concurrency: usize,
    #[clap(long, help = "enable debug mode")]
    debug: bool,
}

fn main() {
    env_logger::builder().filter_level(LevelFilter::Info).init();

    let opts = AppOpts::parse();
    match opts.subcmd {
        SubCommand::Run(opts) => {
            run(&opts).unwrap();
        }
        SubCommand::RunScan(opts) => {
            run_with_scan(&opts).unwrap();
        }
        SubCommand::Scan(opts) => {
            scan_only(&opts).unwrap();
        }
        SubCommand::ShowResults(opts) => {
            show_results(&opts).unwrap();
        }
        SubCommand::ActiveScan(opts) => {
            active_scan(&opts).unwrap();
        }
        SubCommand::Genca(opts) => {
            #[allow(unused_variables)]
            let cert = ca::gen_ca();
            if opts.trust {
                #[cfg(feature = "trust-cert")]
                trust_cert::trust_cert(cert);
            }
        }
    }
}

#[tokio::main]
async fn run(opts: &Run) -> Result<()> {
    info!("CA Private key use: {}", opts.key);
    let private_key_bytes = fs::read(&opts.key).expect("ca private key file path not valid!");
    let private_key = pemfile::pkcs8_private_keys(&mut private_key_bytes.as_slice())
        .expect("Failed to parse private key");
    let private_key = rustls::PrivateKey(private_key[0].clone());

    info!("CA Certificate use: {}", opts.cert);
    let ca_cert_bytes = fs::read(&opts.cert).expect("ca cert file path not valid!");
    let ca_cert =
        pemfile::certs(&mut ca_cert_bytes.as_slice()).expect("Failed to parse CA certificate");
    let ca_cert = rustls::Certificate(ca_cert[0].clone());

    let ca = CertificateAuthority::new(
        private_key,
        ca_cert,
        String::from_utf8(ca_cert_bytes).unwrap(),
        1_000,
    )
    .expect("Failed to create Certificate Authority");

    info!("Http Proxy listen on: http://{}", opts.bind);

    let (rules, mitm_filters) = file::load_rules_amd_mitm_filters(&opts.rule)?;
    let rules = Arc::new(rules);
    let http_handler = RuleHttpHandler::new(rules);

    let proxy = Proxy::builder()
        .ca(ca.clone())
        .listen_addr(opts.bind.parse().expect("bind address not valid!"))
        .upstream_proxy(
            opts.proxy
                .clone()
                .map(|proxy| hyper_proxy::Proxy::new(Intercept::All, proxy.parse().unwrap())),
        )
        .shutdown_signal(shutdown_signal())
        .mitm_filters(mitm_filters.clone())
        .handler(http_handler.clone())
        .build();

    tokio::spawn(proxy.start_proxy());

    tokio::signal::ctrl_c()
        .await
        .expect("failed to listen for event");
    Ok(())
}

#[tokio::main]
async fn run_with_scan(opts: &RunScan) -> Result<()> {
    #[cfg(feature = "passive-scan")]
    {
        use passive_scan::PassiveScanEngine;
        use std::sync::Arc;

        info!("Starting Good-MITM with passive scanning enabled");

        // 创建被动扫描引擎，使用用户提供的配置文件
        let scan_engine = match &opts.scan_config {
            Some(config_path) => {
                info!("Loading scan configuration from: {}", config_path);
                Arc::new(PassiveScanEngine::new())
            },
            None => {
                info!("Using default scan configuration");
                Arc::new(PassiveScanEngine::new_initialized().await)
            }
        };

        info!("Passive scan engine initialized and started");

        // 加载证书
        info!("CA Private key use: {}", opts.key);
        let private_key_bytes = fs::read(&opts.key).expect("ca private key file path not valid!");
        let private_key = pemfile::pkcs8_private_keys(&mut private_key_bytes.as_slice())
            .expect("Failed to parse private key");
        let private_key = rustls::PrivateKey(private_key[0].clone());

        info!("CA Certificate use: {}", opts.cert);
        let ca_cert_bytes = fs::read(&opts.cert).expect("ca cert file path not valid!");
        let ca_cert = pemfile::certs(&mut ca_cert_bytes.as_slice()).expect("Failed to parse CA certificate");
        let ca_cert = rustls::Certificate(ca_cert[0].clone());

        let ca = CertificateAuthority::new(
            private_key,
            ca_cert,
            String::from_utf8(ca_cert_bytes).unwrap(),
            1_000,
        ).expect("Failed to create Certificate Authority");

        info!("Http Proxy listen on: http://{}", opts.bind);

        let (rules, mitm_filters) = file::load_rules_amd_mitm_filters(&opts.rule)?;
        let rules = Arc::new(rules);
        let http_handler = RuleHttpHandler::new(rules);

        // 使用已创建的扫描引擎
        
        // 创建带有被动扫描的HTTP处理器
        let scan_handler = ScanHttpHandler::new(http_handler, scan_engine);

        let proxy = Proxy::builder()
            .ca(ca.clone())
            .listen_addr(opts.bind.parse().expect("bind address not valid!"))
            .upstream_proxy(
                opts.proxy
                    .clone()
                    .map(|proxy| hyper_proxy::Proxy::new(Intercept::All, proxy.parse().unwrap())),
            )
            .shutdown_signal(shutdown_signal())
            .mitm_filters(mitm_filters.clone())
            .handler(scan_handler)
            .build();

        tokio::spawn(proxy.start_proxy());

        info!("Good-MITM with passive scanning is running. Press Ctrl+C to stop.");

        tokio::signal::ctrl_c()
            .await
            .expect("failed to listen for event");

        info!("Shutting down passive scan engine...");
        // PassiveScanEngine不需要显式停止

        Ok(())
    }

    #[cfg(not(feature = "passive-scan"))]
    {
        error!("Passive scanning is not enabled. Please run with --features passive-scan");
        Err(anyhow::anyhow!("Passive scanning feature not enabled"))
    }
}

#[tokio::main]
async fn scan_only(_opts: &Scan) -> Result<()> {
    #[cfg(feature = "passive-scan")]
    {
        use passive_scan::PassiveScanEngine;
        use std::sync::Arc;

        info!("Starting passive scan engine in standalone mode");

        // 移除不存在的ScanConfig引用
        let _scan_config = (); // 占位符

        // 移除对不存在的字段的访问

        // 创建被动扫描引擎并用Arc包装
        let _scan_engine = Arc::new(PassiveScanEngine::new());

        info!("Passive scan engine started successfully");
        info!("Engine is ready to receive events. Use API to submit HTTP events for scanning.");

        // 保持运行直到收到停止信号
        tokio::signal::ctrl_c()
            .await
            .expect("Failed to listen for ctrl-c signal");

        info!("Stopping passive scan engine...");
        // PassiveScanEngine doesn't need explicit stop

        Ok(())
    }

    #[cfg(not(feature = "passive-scan"))]
    {
        error!("Passive scanning is not enabled. Please run with --features passive-scan");
        Err(anyhow::anyhow!("Passive scanning feature not enabled"))
    }
}

fn show_results(opts: &ShowResults) -> Result<()> {
    #[cfg(feature = "passive-scan")]
    {
        // 移除不存在的导入

        info!("Displaying scan results");

        let _storage_path = opts.database.as_deref()
            .unwrap_or("./scan_results");

        // 移除不存在的Storage相关代码
        let _storage = (); // 占位符，但使用正确的名称

        // 使用block_on来在同步函数中调用异步代码
        let _rt = tokio::runtime::Runtime::new()?;
        use passive_scan::ScanResult;
        let results: Vec<ScanResult> = Vec::new(); // 空结果集

        let format = opts.format.as_deref().unwrap_or("table");

        match format {
            "json" => {
                println!("Results: {:?}", results);
            }
            "table" => {
                println!("Scan Results:");
                println!("{}",
                    std::iter::repeat("=").take(80).collect::<String>()
                );

                for result in results {
                    println!(
                        "{:<20} {:<10} {:<30} {}",
                        result.rule_name,
                        format!("{:?}", result.severity),
                        result.url.chars().take(30).collect::<String>(),
                        result.timestamp.format("%Y-%m-%d %H:%M:%S")
                    );
                    println!("  Title: {}", result.title);
                    if !result.description.is_empty() {
                        println!("  Description: {}", result.description.chars().take(80).collect::<String>());
                    }
                    println!();
                }
            }
            _ => {
                error!("Unsupported format: {}", format);
                return Err(anyhow::anyhow!("Unsupported output format"));
            }
        }

        Ok(())
    }

    #[cfg(not(feature = "passive-scan"))]
    {
        error!("Passive scanning is not enabled. Please run with --features passive-scan");
        Err(anyhow::anyhow!("Passive scanning feature not enabled"))
    }
}

#[tokio::main]
async fn active_scan(opts: &ActiveScan) -> Result<()> {
    #[cfg(feature = "active-scan")]
    {
        use active_scan::{ActiveScanConfig, ActiveScanEngine, MatchTarget, RegexMatcher, RegexRule, RuleMatcher, Severity};
        use std::sync::Arc;

        info!("Starting active scan for target: {}", opts.target);

        // 加载或创建配置
        let mut config = if let Some(config_path) = &opts.config {
            info!("Loading active scan configuration from: {}", config_path);
            ActiveScanConfig::from_file(config_path)?
        } else {
            info!("Using default active scan configuration");
            ActiveScanConfig::default()
        };

        // 应用命令行参数覆盖配置
        if let Some(output) = &opts.output {
            config.scan.output_dir = output.clone();
        }
        config.crawler.max_depth = opts.max_depth;
        config.crawler.max_pages = opts.max_pages;
        config.crawler.concurrency = opts.concurrency;

        info!("Active scan configuration:");
        info!("  Target: {}", opts.target);
        info!("  Max depth: {}", config.crawler.max_depth);
        info!("  Max pages: {}", config.crawler.max_pages);
        info!("  Concurrency: {}", config.crawler.concurrency);
        info!("  Output dir: {}", config.scan.output_dir);

        // 创建规则匹配器
        // 这里使用一个简单的示例规则，实际使用时应该从配置文件加载
        let matcher = RegexMatcher::new("active_scan".to_string())
            .add_rule(RegexRule {
                name: "password_exposure".to_string(),
                pattern: regex::Regex::new(r#"(?i)(password|passwd|pwd)\s*[:=]\s*['"]?[\w]{4,}"#).unwrap(),
                severity: Severity::High,
                description: "Potential password exposure detected".to_string(),
                target: MatchTarget::ResponseBody,
                tags: vec!["sensitive".to_string(), "credential".to_string()],
            })
            .add_rule(RegexRule {
                name: "api_key_exposure".to_string(),
                pattern: regex::Regex::new(r#"(?i)(api[_-]?key|apikey)\s*[:=]\s*['"]?[\w]{16,}"#).unwrap(),
                severity: Severity::Critical,
                description: "Potential API key exposure detected".to_string(),
                target: MatchTarget::ResponseBody,
                tags: vec!["sensitive".to_string(), "api_key".to_string()],
            });

        let matcher: Arc<dyn RuleMatcher> = Arc::new(matcher);

        // 创建主动扫描引擎
        let mut engine = ActiveScanEngine::new(config, matcher).await?;

        // 开始扫描
        info!("Starting crawler for target: {}", opts.target);
        let stats = engine.start_scan(&opts.target).await?;

        // 显示统计信息
        info!("Active scan completed!");
        info!("Statistics:");
        info!("  Pages crawled: {}", stats.pages_crawled);
        info!("  Vulnerabilities found: {}", stats.vulnerabilities_found);
        info!("  Errors: {}", stats.errors);

        // 获取并显示结果
        let results = engine.get_results().await?;
        if !results.is_empty() {
            info!("\nVulnerabilities found:");
            for result in &results {
                info!("  [{}] {} - {}",
                    match result.severity {
                        Severity::Critical => "CRITICAL",
                        Severity::High => "HIGH",
                        Severity::Medium => "MEDIUM",
                        Severity::Low => "LOW",
                        Severity::Info => "INFO",
                    },
                    result.title,
                    result.url
                );
            }
        } else {
            info!("No vulnerabilities found.");
        }

        Ok(())
    }

    #[cfg(not(feature = "active-scan"))]
    {
        error!("Active scanning is not enabled. Please build with --features active-scan");
        Err(anyhow::anyhow!("Active scanning feature not enabled"))
    }
}
